{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# `mlc-python`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[MLC-Python](https://github.com/mlc-ai/mlc-python) 是一款以 Python 为核心的工具包，它通过提供 Python 风格的数据类、结构感知工具以及基于 Python 的文本格式，极大地简化了 AI 编译器、运行时环境和复合 AI 系统的开发流程。\n",
    "\n",
    "除了纯粹的 Python 环境外，MLC 还原生支持与 C++ 插件的零拷贝互操作，使得从纯 Python 开发到混合开发乃至无 Python 依赖的工程实践过渡变得顺畅无阻。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 📥 安装\n",
    "\n",
    "```bash\n",
    "pip install -U mlc-python\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 🔑 关键特性\n",
    "\n",
    "### 🏗️ 使用 MLC 数据类定义 IRs\n",
    "\n",
    "MLC提供了 Python 风格的数据类："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import mlc.dataclasses as mlcd\n",
    "\n",
    "@mlcd.py_class(\"demo.MyClass\")\n",
    "class MyClass(mlcd.PyClass):\n",
    "  a: int\n",
    "  b: str\n",
    "  c: float | None\n",
    "\n",
    "instance = MyClass(12, \"test\", c=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "demo.MyClass(a=12, b='test', c=None)"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "instance"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**类型安全**。MLC 数据类通过 Cython 和 C++ 强制执行严格的类型检查。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "demo.MyClass(a=12, b='test', c=10.0)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "instance.c = 10\n",
    "instance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "must be real number, not str",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[4], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43minstance\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mc\u001b[49m \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mwrong type\u001b[39m\u001b[38;5;124m\"\u001b[39m\n",
      "File \u001b[0;32m/media/pc/data/lxw/ai/mlc-python/python/mlc/_cython/base.py:306\u001b[0m, in \u001b[0;36mattach_field.<locals>.fset\u001b[0;34m(this, value, _name)\u001b[0m\n\u001b[1;32m    305\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mfset\u001b[39m(this: typing\u001b[38;5;241m.\u001b[39mAny, value: typing\u001b[38;5;241m.\u001b[39mAny, _name: \u001b[38;5;28mstr\u001b[39m \u001b[38;5;241m=\u001b[39m name) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m--> 306\u001b[0m     \u001b[43msetter\u001b[49m\u001b[43m(\u001b[49m\u001b[43mthis\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mvalue\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32mcore.pyx:1160\u001b[0m, in \u001b[0;36mmlc._cython.core._type_field_accessor.g\u001b[0;34m()\u001b[0m\n",
      "\u001b[0;31mTypeError\u001b[0m: must be real number, not str"
     ]
    }
   ],
   "source": [
    "instance.c = \"wrong type\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "ename": "AttributeError",
     "evalue": "'MyClass' object has no attribute 'non_exist'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[5], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43minstance\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mnon_exist\u001b[49m \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m1\u001b[39m\n",
      "\u001b[0;31mAttributeError\u001b[0m: 'MyClass' object has no attribute 'non_exist'"
     ]
    }
   ],
   "source": [
    "instance.non_exist = 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**序列化**。MLC 数据类支持 {mod}`pickle` 序列化和 JSON 序列化。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "demo.MyClass(a=12, b='test', c=10.0)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "MyClass.from_json(instance.json())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "demo.MyClass(a=12, b='test', c=10.0)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pickle\n",
    "pickle.loads(pickle.dumps(instance))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 🐍 为 IRs 设计基于 Python 的文本格式\n",
    "\n",
    "**打印机**。MLC 查找方法 `__ir_print__` 以将 IR 节点转换为 Python AST：\n",
    "\n",
    "**[[示例](https://github.com/mlc-ai/mlc-python/blob/main/python/mlc/testing/toy_ir/ir.py)]**。将玩具 IR 定义复制到 REPL 中，然后创建下面的 `Func` 节点："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mlc.testing.toy_ir.ir import Var, Func, Assign, Add\n",
    "a, b, c, d, e = Var(\"a\"), Var(\"b\"), Var(\"c\"), Var(\"d\"), Var(\"e\")\n",
    "f = Func(\n",
    "    \"f\", [a, b, c],\n",
    "    stmts=[\n",
    "        Assign(lhs=d, rhs=Add(a, b)),  # d = a + b\n",
    "        Assign(lhs=e, rhs=Add(d, c)),  # e = d + c\n",
    "    ],\n",
    "    ret=e\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "方法 {func}`mlc.printer.to_python` 将 IR 节点转换为基于 Python 的文本；"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "def f(a, b, c):\n",
      "  d = a + b\n",
      "  e = d + c\n",
      "  return e\n"
     ]
    }
   ],
   "source": [
    "import mlc\n",
    "print(mlc.printer.to_python(f))  # 字符串化为Python"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "方法 `mlc.printer.print_python` 进一步以适当的语法高亮渲染文本。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">f</span>(a, b, c):\n",
       "  d <span style=\"color: #AA22FF; font-weight: bold\">=</span> a <span style=\"color: #AA22FF; font-weight: bold\">+</span> b\n",
       "  e <span style=\"color: #AA22FF; font-weight: bold\">=</span> d <span style=\"color: #AA22FF; font-weight: bold\">+</span> c\n",
       "  <span style=\"color: #008000; font-weight: bold\">return</span> e\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "mlc.printer.print_python(f) # 语法高亮"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**AST 解析器**。MLC 提供了一套简洁的 API，用于利用 Python 的 AST 模块实现解析器，包括：\n",
    "- 检查 API，用于获取 Python 类或函数的源代码及其捕获的变量；\n",
    "- 变量管理 API，有助于正确处理作用域；\n",
    "- AST 片段评估 API；\n",
    "- 错误渲染 API。\n",
    "\n",
    "**[[示例](https://github.com/mlc-ai/mlc-python/blob/main/python/mlc/testing/toy_ir/parser.py)]**。借助 MLC API，可以用 100 行代码实现解析器，用于解析上述由 `__ir_printer__` 定义的 Python 文本格式。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 🎯 使用 MLC 结构感知工具测试 IRs\n",
    "\n",
    "通过在 IR 定义中标注 `structure`，MLC 支持结构相等性和结构哈希，以检测 IR 之间的结构等价性：\n",
    "\n",
    "使用 `structure` 定义玩具 IR："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "import mlc.dataclasses as mlcd\n",
    "\n",
    "@mlcd.py_class\n",
    "class Expr(mlcd.PyClass):\n",
    "    def __add__(self, other):\n",
    "        return Add(a=self, b=other)\n",
    "\n",
    "@mlcd.py_class(structure=\"nobind\")\n",
    "class Add(Expr):\n",
    "    a: Expr\n",
    "    b: Expr\n",
    "\n",
    "@mlcd.py_class(structure=\"var\")\n",
    "class Var(Expr):\n",
    "    name: str = mlcd.field(structure=None) # excludes `name` from defined structure\n",
    "\n",
    "@mlcd.py_class(structure=\"bind\")\n",
    "class Let(Expr):\n",
    "    rhs: Expr\n",
    "    lhs: Var = mlcd.field(structure=\"bind\") # `Let.lhs` is the def-site\n",
    "    body: Expr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**结构相等性**。成员方法 `eq_s` 用于比较由MLC的结构化数据类表示的两个IR（中间表示）的结构相等性（即 alpha 等价）。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x, y, z = Var(\"x\"), Var(\"y\"), Var(\"z\")\n",
    "L1 = Let(rhs=x + y, lhs=z, body=z)  # let z = x + y; z\n",
    "L2 = Let(rhs=y + z, lhs=x, body=x)  # let x = y + z; x\n",
    "L3 = Let(rhs=x + x, lhs=z, body=z)  # let z = x + x; z\n",
    "L1.eq_s(L2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [
    {
     "ename": "ValueError",
     "evalue": "Structural equality check failed at {root}.rhs.b: Inconsistent binding. RHS has been bound to a different node while LHS is not bound",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[13], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mL1\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43meq_s\u001b[49m\u001b[43m(\u001b[49m\u001b[43mL3\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43massert_mode\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/media/pc/data/lxw/ai/mlc-python/python/mlc/core/object.py:30\u001b[0m, in \u001b[0;36mObject.eq_s\u001b[0;34m(self, other, bind_free_vars, assert_mode)\u001b[0m\n\u001b[1;32m     23\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21meq_s\u001b[39m(\n\u001b[1;32m     24\u001b[0m     \u001b[38;5;28mself\u001b[39m,\n\u001b[1;32m     25\u001b[0m     other: Object,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     28\u001b[0m     assert_mode: \u001b[38;5;28mbool\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m,\n\u001b[1;32m     29\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m \u001b[38;5;28mbool\u001b[39m:\n\u001b[0;32m---> 30\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mPyAny\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_mlc_eq_s\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mother\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mbind_free_vars\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43massert_mode\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32mcore.pyx:339\u001b[0m, in \u001b[0;36mmlc._cython.core.PyAny._mlc_eq_s\u001b[0;34m()\u001b[0m\n",
      "File \u001b[0;32mcore.pyx:1219\u001b[0m, in \u001b[0;36mmlc._cython.core.func_call\u001b[0;34m()\u001b[0m\n",
      "File \u001b[0;32mcore.pyx:653\u001b[0m, in \u001b[0;36mmlc._cython.core._func_call_impl\u001b[0;34m()\u001b[0m\n",
      "File \u001b[0;32mcore.pyx:644\u001b[0m, in \u001b[0;36mmlc._cython.core._func_call_impl_with_c_args\u001b[0;34m()\u001b[0m\n",
      "File \u001b[0;32mcore.pyx:286\u001b[0m, in \u001b[0;36mmlc._cython.core._check_error_from\u001b[0;34m()\u001b[0m\n",
      "File \u001b[0;32m/media/pc/data/lxw/ai/mlc-python/include/mlc/core/./func.h:30\u001b[0m, in \u001b[0;36mmlc::FuncObj::SafeCallImpl(mlc::FuncObj const*, int, mlc::AnyView const*, mlc::Any*)\u001b[0;34m()\u001b[0m\n\u001b[1;32m     28\u001b[0m static int32_t SafeCallImpl(const FuncObj *self, int32_t num_args, const AnyView *args, Any *ret) {\n\u001b[1;32m     29\u001b[0m   MLC_SAFE_CALL_BEGIN();\n\u001b[0;32m---> 30\u001b[0m   self->call(self, num_args, args, ret);\n\u001b[1;32m     31\u001b[0m   MLC_SAFE_CALL_END(ret);\n\u001b[1;32m     32\u001b[0m }\n",
      "File \u001b[0;32m/media/pc/data/lxw/ai/mlc-python/include/mlc/core/./func_details.h:145\u001b[0m, in \u001b[0;36mvoid mlc::core::FuncCallUnpacked<bool (*)(mlc::Object*, mlc::Object*, bool, bool)>(mlc::FuncObj const*, int, mlc::AnyView const*, mlc::Any*)\u001b[0;34m()\u001b[0m\n\u001b[1;32m    143\u001b[0m   using IdxSeq = std::make_index_sequence<N>;\n\u001b[1;32m    144\u001b[0m   using RetType = typename FuncCanonicalize<FuncType>::RetType;\n\u001b[0;32m--> 145\u001b[0m   UnpackCall<RetType, typename Traits::ArgType>::template Run<FuncType>(\n\u001b[1;32m    146\u001b[0m       &static_cast<const FuncImpl<FuncType> *>(obj)->func_, args, ret, IdxSeq{});\n\u001b[1;32m    147\u001b[0m }\n",
      "File \u001b[0;32m/media/pc/data/lxw/ai/mlc-python/include/mlc/core/./func_details.h:123\u001b[0m, in \u001b[0;36mvoid mlc::core::UnpackCall<bool, std::tuple<mlc::Object*, mlc::Object*, bool, bool> >::Run<bool (*)(mlc::Object*, mlc::Object*, bool, bool), bool (*)(mlc::Object*, mlc::Object*, bool, bool), 0ul, 1ul, 2ul, 3ul>(bool (**)(mlc::Object*, mlc::Object*, bool, bool), mlc::AnyView const*, mlc::Any*, std::integer_sequence<unsigned long, 0ul, 1ul, 2ul, 3ul>)\u001b[0;34m()\u001b[0m\n\u001b[1;32m    121\u001b[0m     (*func)(CVT::template AsType<Args, I>::Run(args[I], nullptr)...);\n\u001b[1;32m    122\u001b[0m   } else if constexpr (Storage::total == 0 && !std::is_void_v<RetType>) {\n\u001b[0;32m--> 123\u001b[0m     *ret = (*func)(CVT::template AsType<Args, I>::Run(args[I], nullptr)...);\n\u001b[1;32m    124\u001b[0m   }\n\u001b[1;32m    125\u001b[0m }\n",
      "File \u001b[0;32m/media/pc/data/lxw/ai/mlc-python/cpp/structure.cc:35\u001b[0m, in \u001b[0;36mStructuralEqual\u001b[0;34m()\u001b[0m\n\u001b[1;32m     33\u001b[0m     std::ostringstream os;\n\u001b[1;32m     34\u001b[0m     os << \"Structural equality check failed at \" << e.path << \": \" << e.what();\n\u001b[0;32m---> 35\u001b[0m     MLC_THROW(ValueError) << os.str();\n\u001b[1;32m     36\u001b[0m   }\n\u001b[1;32m     37\u001b[0m }\n",
      "\u001b[0;31mValueError\u001b[0m: Structural equality check failed at {root}.rhs.b: Inconsistent binding. RHS has been bound to a different node while LHS is not bound"
     ]
    }
   ],
   "source": [
    "L1.eq_s(L3, assert_mode=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**结构哈希**。MLC 数据类的结构可以通过 `hash_s` 进行哈希计算，这保证了如果两个数据类是 alpha 等价的，它们将具有相同的结构哈希值："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "L1_hash, L2_hash, L3_hash = L1.hash_s(), L2.hash_s(), L3.hash_s()\n",
    "assert L1_hash == L2_hash\n",
    "assert L1_hash != L3_hash"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ⚡ 逐步迁移到 C++ 并使用 MLC 插件\n",
    "\n",
    "(🚧 正在建设中)\n",
    "\n",
    "MLC 无缝支持与 C++ 插件的零拷贝双向互操作性，且无需额外依赖。通过逐步迁移类和方法，可以将纯 Python 原型过渡到混合开发或无 Python 开发模式。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ⛽ 开发\n",
    "\n",
    "### ⚙️ 可编辑构建\n",
    "\n",
    "```bash\n",
    "pip install --verbose --editable \".[dev]\"\n",
    "pre-commit install\n",
    "```\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 🎡 构建 Wheels\n",
    "\n",
    "本项目使用 `cibuildwheel` 来构建跨平台的 wheels。更多详情请参阅 `.github/workflows/wheels.yml`。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ai",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
